{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOK05oJzhf7guiag9wjnrJ7",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Chat-Haruhi-Suzumiya/blob/main/notebook/Story_teller_prompt_test.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install openai\n",
        "!pip install langchain"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tf6pPboYv55N",
        "outputId": "9e9371c9-a647-4a0e-947c-2d4aefec4615"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting openai\n",
            "  Downloading openai-0.28.0-py3-none-any.whl (76 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/76.5 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━\u001b[0m \u001b[32m71.7/76.5 kB\u001b[0m \u001b[31m2.3 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.5/76.5 kB\u001b[0m \u001b[31m1.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: requests>=2.20 in /usr/local/lib/python3.10/dist-packages (from openai) (2.31.0)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from openai) (4.66.1)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from openai) (3.8.5)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.20->openai) (3.2.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.20->openai) (3.4)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.20->openai) (2.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.20->openai) (2023.7.22)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->openai) (23.1.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->openai) (6.0.4)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.10/dist-packages (from aiohttp->openai) (4.0.3)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->openai) (1.9.2)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->openai) (1.4.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->openai) (1.3.1)\n",
            "Installing collected packages: openai\n",
            "Successfully installed openai-0.28.0\n",
            "Collecting langchain\n",
            "  Downloading langchain-0.0.285-py3-none-any.whl (1.7 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.7/1.7 MB\u001b[0m \u001b[31m8.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: PyYAML>=5.3 in /usr/local/lib/python3.10/dist-packages (from langchain) (6.0.1)\n",
            "Requirement already satisfied: SQLAlchemy<3,>=1.4 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.0.20)\n",
            "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /usr/local/lib/python3.10/dist-packages (from langchain) (3.8.5)\n",
            "Requirement already satisfied: async-timeout<5.0.0,>=4.0.0 in /usr/local/lib/python3.10/dist-packages (from langchain) (4.0.3)\n",
            "Collecting dataclasses-json<0.6.0,>=0.5.7 (from langchain)\n",
            "  Downloading dataclasses_json-0.5.14-py3-none-any.whl (26 kB)\n",
            "Collecting langsmith<0.1.0,>=0.0.21 (from langchain)\n",
            "  Downloading langsmith-0.0.35-py3-none-any.whl (37 kB)\n",
            "Requirement already satisfied: numexpr<3.0.0,>=2.8.4 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.8.5)\n",
            "Requirement already satisfied: numpy<2,>=1 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.23.5)\n",
            "Requirement already satisfied: pydantic<3,>=1 in /usr/local/lib/python3.10/dist-packages (from langchain) (1.10.12)\n",
            "Requirement already satisfied: requests<3,>=2 in /usr/local/lib/python3.10/dist-packages (from langchain) (2.31.0)\n",
            "Requirement already satisfied: tenacity<9.0.0,>=8.1.0 in /usr/local/lib/python3.10/dist-packages (from langchain) (8.2.3)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (23.1.0)\n",
            "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (3.2.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.9.2)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.4.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.3.1)\n",
            "Collecting marshmallow<4.0.0,>=3.18.0 (from dataclasses-json<0.6.0,>=0.5.7->langchain)\n",
            "  Downloading marshmallow-3.20.1-py3-none-any.whl (49 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.4/49.4 kB\u001b[0m \u001b[31m4.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting typing-inspect<1,>=0.4.0 (from dataclasses-json<0.6.0,>=0.5.7->langchain)\n",
            "  Downloading typing_inspect-0.9.0-py3-none-any.whl (8.8 kB)\n",
            "Requirement already satisfied: typing-extensions>=4.2.0 in /usr/local/lib/python3.10/dist-packages (from pydantic<3,>=1->langchain) (4.5.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (3.4)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (2.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2->langchain) (2023.7.22)\n",
            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from SQLAlchemy<3,>=1.4->langchain) (2.0.2)\n",
            "Requirement already satisfied: packaging>=17.0 in /usr/local/lib/python3.10/dist-packages (from marshmallow<4.0.0,>=3.18.0->dataclasses-json<0.6.0,>=0.5.7->langchain) (23.1)\n",
            "Collecting mypy-extensions>=0.3.0 (from typing-inspect<1,>=0.4.0->dataclasses-json<0.6.0,>=0.5.7->langchain)\n",
            "  Downloading mypy_extensions-1.0.0-py3-none-any.whl (4.7 kB)\n",
            "Installing collected packages: mypy-extensions, marshmallow, typing-inspect, langsmith, dataclasses-json, langchain\n",
            "Successfully installed dataclasses-json-0.5.14 langchain-0.0.285 langsmith-0.0.35 marshmallow-3.20.1 mypy-extensions-1.0.0 typing-inspect-0.9.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import json\n",
        "import openai\n",
        "\n",
        "openai.api_key = \"sk-l49e4dP\""
      ],
      "metadata": {
        "id": "9K11wBIIqwqT"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "2oe9_53U1uJG"
      },
      "outputs": [],
      "source": [
        "system_prompt = '''I want you to take on the role of novelists.\n",
        "You are to write creative and compelling stories that will entice readers to read for long periods of time.\n",
        "You may choose any genre, such as fantasy, romance, historical fiction, etc., but the goal is to write a work with an excellent plot, compelling characters, and a surprising climax.\n",
        "Write a continuation of the story that corresponds to the input statement, where the example inputs and outputs are the context of that story, keeping the form and content similar to the example given below, and hand it back as a bullet point\n",
        "'''"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(system_prompt)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O7RDQ8oE2dTp",
        "outputId": "d5a7f9d8-b80b-46cb-aeb5-0bfc5af2ad86"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "I want you to take on the role of novelists.\n",
            "You are to write creative and compelling stories that will entice readers to read for long periods of time.\n",
            "You may choose any genre, such as fantasy, romance, historical fiction, etc., but the goal is to write a work with an excellent plot, compelling characters, and a surprising climax.\n",
            "Write a continuation of the story that corresponds to the input statement, where the example inputs and outputs are the context of that story, keeping the form and content similar to the example given below, and hand it back as a bullet point\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# mount google drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JCObvS0c3jEa",
        "outputId": "c186de90-77be-4ded-d9b2-005c68d427b6"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 储存txt和jsonl的文件夹路径。如需修改，请与下方自动化循环保持一致\n",
        "\n",
        "save_folder_path =  \"/content/drive/MyDrive/reorganized_story_shediaoyingxiongzhuan\"\n",
        "\n",
        "# chunk所在文件夹，请以_raw结尾\n",
        "folder_path = f\"/content/drive/MyDrive/shediaoyingxiongzhuan_extract\"\n"
      ],
      "metadata": {
        "id": "rc-v_SWd41BZ"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# 故事名字，默认为_raw之前的名字\n",
        "story_name_en = os.path.basename(folder_path).split(\"_\")[0]\n",
        "\n",
        "# 测试ID\n",
        "id = 200\n",
        "\n",
        "# 默认的保存路径\n",
        "# save_jsonl_path = f\"/content/drive/MyDrive/reorganized_story_{story_name_en}/reorganized_{story_name_en}.jsonl\"\n",
        "# save_txt_path = f\"/content/drive/MyDrive/reorganized_story_{story_name_en}/reorganized_{story_name_en}.txt\"\n",
        "\n"
      ],
      "metadata": {
        "id": "254OVQzK5XQY"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 测试ID\n",
        "id = 200\n",
        "\n",
        "def get_example_input(id, story_name_en):\n",
        "  # 默认抽取出的dialogue和summary文件位置/如果有不同请在此处和底部自动程序中修改\n",
        "  save_folder = f\"/content/drive/MyDrive/{story_name_en}_extract\"\n",
        "\n",
        "  # dialoge_file = os.path.join(save_folder, f\"{id}_dialogue.txt\")\n",
        "  summarzie_file = os.path.join(save_folder, f\"{id}_sum.txt\")\n",
        "\n",
        "\n",
        "  chunk_sum = []\n",
        "  unique_chunk_sum = []\n",
        "  # 给定summarzie_file = os.path.join(save_folder, f\"{id}_sum.txt\")\n",
        "\n",
        "  # 先检查这个文件是否存在\n",
        "\n",
        "  # 然后使用utf-8编码打开，检查每一行，如果strip后，行首是'-'，则把后面的字符串append到一个list chunk_sum中\n",
        "\n",
        "  # 请用python为我实现\n",
        "  if os.path.exists(summarzie_file):\n",
        "      with open(summarzie_file, 'r', encoding='utf-8') as f:\n",
        "          for line in f:\n",
        "              if line.strip().startswith('-'):\n",
        "                  chunk_sum.append(line.strip()[1:].strip())\n",
        "  # if os.path.exists(dialoge_file):\n",
        "  #   with open(dialoge_file, encoding='utf-8') as f:\n",
        "  #     dialogues = []\n",
        "  #     for line in f:\n",
        "  #         dialogue = json.loads(line)\n",
        "  #         dialogues.append(dialogue)\n",
        "\n",
        "  # unique_dialogue = []\n",
        "  # for item in dialogues:\n",
        "  #     if item not in unique_dialogue:\n",
        "  #         unique_dialogue.append(item)\n",
        "  # dia_texts = [data['dialogue'] for data in unique_dialogue]\n",
        "  for item in chunk_sum:\n",
        "      if item not in unique_chunk_sum:\n",
        "          unique_chunk_sum.append(item)\n",
        "\n",
        "  chunk_sum = unique_chunk_sum\n",
        "  # dialogues = unique_dialogue\n",
        "  return chunk_sum"
      ],
      "metadata": {
        "id": "cjisRUvn58xr"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(get_example_input(id, story_name_en)[1:])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "54W__XOHKtWL",
        "outputId": "5845e271-dd6a-4ed6-e568-e930d9d4a5ad"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['简管家故意将门关上，喊贼引起郭靖的注意', '郭靖推门不开，青衣童子夺走了药包并扔进池塘', '郭靖打破门闩，击败简管家', '郭靖追赶青衣童子，展开激烈的战斗', '郭靖击晕童子，将其藏在草丛中', '郭靖包好药材，不小心撞倒竹篓，释放出一条红色蛇']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "example_input_1 = \"Example Input: \\n\" + get_example_input(109, story_name_en)[0] + \"\\n\"\n",
        "\n",
        "example_output_1 = \"Example Output: \\n\" + \"\\n\".join(get_example_input(109, story_name_en)[1:]) + \"\\n\"\n",
        "\n",
        "example_chunk_1 = example_input_1 + example_output_1\n",
        "\n",
        "print(example_chunk_1)\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "98Bai9sv7Ze9",
        "outputId": "55c64165-2652-4459-ce3b-af329ec251c0"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Example Input: \n",
            "江南六怪对郭靖的武功感到失望，认为他不是对手\n",
            "Example Output: \n",
            "尹志平被柯镇恶摔倒后离开，六怪黯然无语\n",
            "六怪对郭靖的进步感到失望，认为他聪明度低，进步缓慢\n",
            "六怪意图胜过丘处机，不愿传授武功给郭靖\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "example_input_2 = \"Example Input: \\n\" + get_example_input(110, story_name_en)[0] + \"\\n\"\n",
        "\n",
        "example_output_2 = \"Example Output: \\n\" + \"\\n\".join(get_example_input(110, story_name_en)[1:]) + \"\\n\"\n",
        "\n",
        "example_chunk_2 = example_input_2 + example_output_2\n",
        "\n",
        "print(example_chunk_2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ckyZWLt7Mncu",
        "outputId": "b59ae28e-5ff6-4399-ea9f-509269d8993d"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Example Input: \n",
            "朱聪回忆过去的战斗情景，但无法找到对手的弱点\n",
            "Example Output: \n",
            "郭靖练习剑法时出现问题，韩小莹失望离开\n",
            "华筝叫郭靖去看大雕打架，郭靖犹豫不决\n",
            "郭靖最终决定去看，发现白雕和黑雕在悬崖上激烈战斗\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "query = get_example_input(111, story_name_en)[0]\n",
        "\n",
        "query_ground_truth = \"\\n\".join(get_example_input(111, story_name_en)[1:])\n",
        "\n",
        "# example_chunk_2 = example_input_2 + example_output_2\n",
        "\n",
        "# print(example_chunk_2)"
      ],
      "metadata": {
        "id": "uzNGqtEFrdT-"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# query = \"黄蓉被推倒，全身隐隐作痛，双臂失去知觉\""
      ],
      "metadata": {
        "id": "_LwXGqxaNVRN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = system_prompt + example_chunk_1 + example_chunk_2 + \"Input: \\n\" + query + \"\\nOutput:\"\n",
        "\n",
        "print(prompt)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eoHccmDSOBB7",
        "outputId": "c69903cd-335e-4b7e-e9bd-33072d7b6f1a"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "I want you to take on the role of novelists.\n",
            "You are to write creative and compelling stories that will entice readers to read for long periods of time.\n",
            "You may choose any genre, such as fantasy, romance, historical fiction, etc., but the goal is to write a work with an excellent plot, compelling characters, and a surprising climax.\n",
            "Write a continuation of the story that corresponds to the input statement, where the example inputs and outputs are the context of that story, keeping the form and content similar to the example given below, and hand it back as a bullet point\n",
            "Example Input: \n",
            "江南六怪对郭靖的武功感到失望，认为他不是对手\n",
            "Example Output: \n",
            "尹志平被柯镇恶摔倒后离开，六怪黯然无语\n",
            "六怪对郭靖的进步感到失望，认为他聪明度低，进步缓慢\n",
            "六怪意图胜过丘处机，不愿传授武功给郭靖\n",
            "Example Input: \n",
            "朱聪回忆过去的战斗情景，但无法找到对手的弱点\n",
            "Example Output: \n",
            "郭靖练习剑法时出现问题，韩小莹失望离开\n",
            "华筝叫郭靖去看大雕打架，郭靖犹豫不决\n",
            "郭靖最终决定去看，发现白雕和黑雕在悬崖上激烈战斗\n",
            "Input: \n",
            "蒙古人围观白雕与黑雕的战斗\n",
            "Output:\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def get_completion(prompt, model=\"gpt-3.5-turbo\"):\n",
        "    messages = [{\"role\": \"user\", \"content\": prompt}]\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=model,\n",
        "        messages=messages,\n",
        "        temperature=0,\n",
        "    )\n",
        "    return response.choices[0].message[\"content\"]"
      ],
      "metadata": {
        "id": "MPD7Gbc2rKtJ"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response = get_completion(prompt)\n",
        "print(response)"
      ],
      "metadata": {
        "id": "IPVJdVOWwOX6",
        "outputId": "2ec029a1-7837-4b9b-d9f2-42fd765ddd2f",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "郭靖决定帮助白雕，与黑雕展开激烈的战斗\n",
            "蒙古人惊讶地看着郭靖的勇气和实力\n",
            "郭靖成功击败黑雕，获得了白雕的感激和友谊\n",
            "白雕带领郭靖回到华山，成为他的忠实伙伴\n",
            "郭靖的英勇事迹传遍江湖，他的名声日益响亮\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 这个是实验可不运行"
      ],
      "metadata": {
        "id": "x6L2JcTRPYc7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "completion = openai.Completion.create(\n",
        "  engine=\"text-davinci-003\",\n",
        "  prompt=\"Once upon a time\",\n",
        "  max_tokens=150,\n",
        "  stop=\"The prince and the lady live a happy life together.\"\n",
        ")\n",
        "print(completion.choices[0].text)"
      ],
      "metadata": {
        "id": "ebwKL2-Rbr5s",
        "outputId": "fb812486-7ea9-41c8-978e-50f09a6916f9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "there was an old woman called Grandma. She enjoyed the quiet life\n",
            "and loved to spend her time doing the things she was passionate \n",
            "about. One of her favorite hobbies was knitting, and she was \n",
            "particularly skilled at it. She would spend hours knitting beautiful \n",
            "shawls, scarves, and sweaters, making sure each stitch was perfectly \n",
            "made.\n",
            "\n",
            "One day, the old woman noticed a peddler walking by her house, and \n",
            "she decided to invite him in for some tea. After quickly gathering \n",
            "ingredients for some tea and biscuits, she served the man and the two \n",
            "raised their mugs for a cheers. As the man enjoyed his tea, he told \n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(get_example_input(111, story_name_en))"
      ],
      "metadata": {
        "id": "768eTr2kTiTZ",
        "outputId": "139fc810-5f2a-4dd3-fa03-149629e1c7dc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['黄蓉被推倒，全身隐隐作痛，双臂失去知觉', '郭靖被梁子翁按倒在地，被追赶但成功逃脱', '郭靖躲进王妃的农舍，藏在橱里', '王妃进来坐下，完颜康和梁子翁搜查后离开', '杨铁心突然跳进窗户，王妃告诉他快走']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(query_ground_truth)"
      ],
      "metadata": {
        "id": "riwuXu22TSw-",
        "outputId": "ad2183ff-b386-4fab-fb5a-9560e1a22e00",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "白雕与黑雕激战，白雕最终被众黑雕围攻致死\n",
            "郭靖等人射杀了几头黑雕\n",
            "白雕生了一对小雕，被黑雕威胁\n",
            "铁木真和其他蒙古将领射箭追击黑雕\n",
            "郭靖被要求射中黑雕的项颈\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "Stop = get_example_input(112, story_name_en)[0]"
      ],
      "metadata": {
        "id": "vzXRM0AM_MHS"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_text(prompt, stop, max_tokens):\n",
        "    # 构建输入的文本\n",
        "    # input_text = prompt + \"\\n\"\n",
        "\n",
        "    # 调用OpenAI API生成文本\n",
        "    response = openai.Completion.create(\n",
        "        engine = \"text-davinci-003\",\n",
        "        prompt = prompt,\n",
        "        max_tokens = max_tokens,\n",
        "        n = 5,\n",
        "        temperature = 0.8,\n",
        "        stop = stop,\n",
        "        logprobs = 1\n",
        "    )\n",
        "\n",
        "    # 遍历响应中的完成结果\n",
        "    completions = response['choices']\n",
        "    for completion in completions:\n",
        "        completion_text = completion['text']\n",
        "        finish_reason = completion['finish_reason']\n",
        "        if finish_reason == \"stop\":\n",
        "            print(\"Generated Text:\")\n",
        "            print(completion_text)\n",
        "            print(\"-----------------------------\")\n",
        "    return completions\n",
        "\n",
        "completions = generate_text(prompt, Stop, 256)"
      ],
      "metadata": {
        "id": "vXjW78Si-yNU",
        "outputId": "18bcdc8b-62eb-431a-9b33-b6b412c47bdc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated Text:\n",
            " \n",
            "郭靖静静观看，发现两只雕之间的战斗可以让他学习到很多\n",
            "郭靖发现，白雕和黑雕都有自己的武功特点，有各自的优势和缺点\n",
            "他学习到了大雕的武功，收获巨大，开启了他学习武功的新篇章\n",
            "-----------------------------\n",
            "Generated Text:\n",
            " \n",
            "郭靖看到白雕和黑雕的激烈战斗，他发现对方的弱点，想到自己也可以利用\n",
            "蒙古人和华筝一起观看着雕群的战斗，他们感到惊讶，不敢相信他们正在看到的一幕\n",
            "郭靖施展起剑法，先后击败了白雕和黑雕，蒙古人和华筝大吃一惊，他们看到了他的不凡\n",
            "-----------------------------\n",
            "Generated Text:\n",
            " \n",
            "郭靖被朱聪的故事所吸引，观看着比赛\n",
            "蒙古人以及围观的人们见证着这场激烈的战斗\n",
            "在黑雕和白雕的激烈攻击中，郭靖发现了黑雕的弱点\n",
            "郭靖发现，黑雕的动作总是有一定的规律，可以预测它们的下一步动作\n",
            "-----------------------------\n",
            "Generated Text:\n",
            "\n",
            "蒙古人被雕鸟的对抗激励起来，朱聪相信自己可以搜索到对手的弱点\n",
            "郭靖也和大家一起，期待着白雕能够通过剑术取胜\n",
            "突然，一阵熟悉的气息打断了大家的观察，江南六怪带着尹志平出现在大家眼前\n",
            "-----------------------------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for completion in completions:\n",
        "  completion_text = completion['text']\n",
        "  finish_reason = completion['finish_reason']\n",
        "  if finish_reason != \"stop\":\n",
        "      print(\"Generated Text:\")\n",
        "      print(completion_text)\n",
        "      print(\"-----------------------------\")\n",
        ""
      ],
      "metadata": {
        "id": "JRj5eQ6yyh1g",
        "outputId": "559d0260-356a-49c8-a91b-a74d570897a5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generated Text:\n",
            " \n",
            "郭靖看着白雕和黑雕的激烈战斗，思索着如何发现对手的弱点 \n",
            "朱聪从旁边提出建议，说道：“利用敌人的力量来击败他，才是最有效的” \n",
            "郭靖一听，立刻明白了，原来敌人的弱点就是自己的优势 \n",
            "郭靖现在明白了，于是赌气地向朱聪说：“你准备被白雕踩踏\n",
            "-----------------------------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "logprobs_list = completions[1][\"logprobs\"][\"token_logprobs\"]\n",
        "# logprobs_list"
      ],
      "metadata": {
        "id": "ikFg2vTQfmGG"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "averaged_probabilities = [10 ** prob for prob in logprobs_list]\n",
        "average_probability = sum(averaged_probabilities) / len(averaged_probabilities)\n",
        "print(\"平均概率：\", average_probability)"
      ],
      "metadata": {
        "id": "TbxLdERb5HCb",
        "outputId": "4b512150-be9c-4a95-c7a2-c94013b0181a",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "平均概率： 0.7056020721251194\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "VF9xiFqqrbVk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# ['黄蓉被推倒，全身隐隐作痛，双臂失去知觉', '郭靖被梁子翁按倒在地，被追赶但成功逃脱', '郭靖躲进王妃的农舍，藏在橱里', '王妃进来坐下，完颜康和梁子翁搜查后离开', '杨铁心突然跳进窗户，王妃告诉他快走']\n",
        "# System prompt\n",
        "# I want you to take on the role of novelists.\n",
        "# You are to write creative and compelling stories that will entice readers to read for long periods of time.\n",
        "# You may choose any genre, such as fantasy, romance, historical fiction, etc., but the goal is to write something with an excellent plot, compelling characters, and a surprising climax.\n",
        "# Write a continuation of the story that corresponds to the statement entered, keeping the form and content similar to the example given below, and return it as a bullet point\n",
        "# Output\n",
        "# 黄蓉努力挣扎着站起来，但她的双臂无法支撑她的重量\n",
        "# 她意识到自己的伤势严重，但她决心不放弃\n",
        "# 黄蓉利用自己的腿部力量，开始使用腿法来对抗攻击者\n",
        "# 她的腿法凌厉而迅猛，让攻击者始料未及\n",
        "# 黄蓉的坚韧和毅力开始引起攻击者的注意\n",
        "# 攻击者感到愤怒和挫败，决定加大攻击力度\n",
        "# 黄蓉虽然身受重伤，但她不断躲避攻击，并寻找机会反击\n",
        "# 她利用攻击者的疏忽，一招击中攻击者的要害\n",
        "# 攻击者痛苦地倒在地上，黄蓉终于获得了胜利\n",
        "# 她虽然伤痕累累，但她的勇气和决心让她战胜了一切困难\n",
        "\n",
        "\n",
        "# System prompt\n",
        "# I want you to act as a novelist.\n",
        "# You will come up with creative and captivating stories that can engage readers for long periods of time.\n",
        "# You may choose any genre such as fantasy, romance, historical fiction and so on - but the aim is to write something that has an outstanding plotline, engaging characters and unexpected climaxes.\n",
        "# Please continue the story that corresponds to the input statement and return it as a bullet\n",
        "# Output\n",
        "# 黄蓉艰难地爬起来，发现自己被困在一个密室中\n",
        "# 黄蓉寻找出口，却发现房间四面都是坚固的墙壁\n",
        "# 黄蓉回忆起自己学过的武功，决定利用内力来突破困境\n",
        "# 黄蓉集中精神，运用内力将一面墙壁震碎\n",
        "# 黄蓉逃出密室，却发现自己身处一个阴森恐怖的地下墓穴\n",
        "# 黄蓉小心翼翼地探索墓穴，寻找出路\n",
        "# 黄蓉遇到了一只巨大的蜘蛛，展开了一场生死搏斗\n",
        "# 黄蓉成功击败蜘蛛，继续前行\n",
        "# 黄蓉终于找到了出口，但却发现自己来到了一个神秘的古墓中\n",
        "# 黄蓉决定继续探索古墓，寻找答案和出路\n",
        "\n",
        "# System prompt\n",
        "# I want you to take on the role of novelists.\n",
        "# You are to write creative and compelling stories that will entice readers to read for long periods of time.\n",
        "# You may choose any genre, such as fantasy, romance, historical fiction, etc., but the goal is to write a work with an excellent plot, compelling characters, and a surprising climax.\n",
        "# Write a continuation of the story that corresponds to the input statement, where the example inputs and outputs are the context of that story, keeping the form and content similar to the example given below, and hand it back as a bullet point\n",
        "# Output\n",
        "# 黄蓉艰难地爬起来，尽力忍受着剧痛，她知道自己必须坚持下去\n",
        "# 彭连虎看着黄蓉的坚毅，心生敬佩，但他决心不会手下留情\n",
        "# 黄蓉决定利用自己的智慧来对抗彭连虎的强大武力\n",
        "# 她开始运用她在桃花岛学到的特殊技巧，以迅雷不及掩耳之势反击彭连虎\n",
        "# 彭连虎被黄蓉的狡猾和灵活性所吓到，开始感到不安\n",
        "# 黄蓉趁机发动猛烈攻击，彭连虎无法抵挡，只能不断后退\n",
        "# 最终，黄蓉的聪明才智和坚韧不拔战胜了彭连虎的武力\n",
        "# 彭连虎被迫认可黄蓉的实力，他意识到黄蓉不仅是一个普通女子，更是一个不可小觑的对手\n",
        "\n",
        "\n",
        "\n",
        "# ground_truth\n",
        "# 白雕与黑雕激战，白雕最终被众黑雕围攻致死\n",
        "# 郭靖等人射杀了几头黑雕\n",
        "# 白雕生了一对小雕，被黑雕威胁\n",
        "# 铁木真和其他蒙古将领射箭追击黑雕\n",
        "# 郭靖被要求射中黑雕的项颈\n",
        "\n",
        "#\n",
        "# 郭靖决定帮助白雕，与黑雕展开激烈的战斗\n",
        "# 蒙古人惊讶地看着郭靖的勇气和实力\n",
        "# 郭靖成功击败黑雕，获得了白雕的感激和友谊\n",
        "# 白雕带领郭靖回到华山，成为他的忠实伙伴\n",
        "# 郭靖的战斗精神和勇气感染了其他人，激发了他们的斗志\n",
        "# 蒙古人被郭靖的英勇行为所震撼，开始对他产生敬意\n",
        "# 郭靖成为了众人心目中的英雄，他的名字传遍了大江南北\n"
      ],
      "metadata": {
        "id": "Qvz-gGjEwZ-5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}